# -*- coding: utf-8 -*-
"""
文件名称: batch_eval_steps.py

功能概述:
    1. 批量遍历某个评估输出目录(例如 experiments/eval_output/sub_exp_1_p20),
       自动发现所有形如 "eval_step_XXX.jsonl" 的文件 (XXX 可转为 int 的 step)。
    2. 对每一个 step 的 JSONL 调用已有的单文件评估脚本 cal_score.py, 得到:
         - 打分后的 JSONL (每题 score + reason)
         - 对应的 metrics JSON (overall 及各题型 strict_accuracy 等)
    3. 从每个 step 的 metrics JSON 中提取以下 5 个 strict_accuracy:
         - overall_strict_accuracy
         - judge_strict_accuracy
         - compare_strict_accuracy
         - diagnosis_strict_accuracy
         - choice_strict_accuracy
       并在内存中构建 step -> metrics 的映射。
    4. 将所有 step 的指标汇总写入:
         - <exp_dir>/scored/step_metrics_summary.json
    5. 按 step 升序绘制 5 条折线图 (灰色背景 + 网格 + 方形点 + 深色线条),
       保存为 <exp_dir>/scored/step_metrics_plot.png。

使用示例:
    python fundus_reasoner/eval/batch_eval_steps.py \
        --exp-dir experiments/eval_output/sub_exp_1_p20 \
        --cal-score-script fundus_reasoner/eval/cal_score.py

注意事项:
    - 不使用并发, 顺序依次调用 cal_score.py, 避免并发干扰与调试困难。
    - 所有文件 IO 与子进程调用出错时均直接 raise RuntimeError, 并包含关键调试信息。
"""

import argparse
import json
import logging
import re
import subprocess
import sys
from pathlib import Path
from typing import Dict, Any, List, Tuple
import os
import matplotlib.pyplot as plt


logger = logging.getLogger(__name__)
if not logger.handlers:
    handler = logging.StreamHandler()
    formatter = logging.Formatter(
        fmt="[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)


# -----------------------------
# 工具函数: 发现 step -> 文件 映射
# -----------------------------
def discover_step_files(exp_dir: Path) -> Dict[int, Path]:
    """在指定目录下发现所有 eval_step_XXX.jsonl 文件

    规则:
        - 文件名必须形如: eval_step_<int>.jsonl
        - 中间的部分必须能转为 int
    返回:
        {step: file_path}
        若存在重复 step, 则抛出异常。
    """
    if not exp_dir.is_dir():
        raise RuntimeError(f"评估目录不存在或不是目录: {exp_dir}")

    pattern = re.compile(r"^eval_step_(\d+)\.jsonl$")

    step_to_file: Dict[int, Path] = {}

    for f in exp_dir.iterdir():
        if not f.is_file():
            continue
        m = pattern.match(f.name)
        if not m:
            continue
        step_str = m.group(1)
        try:
            step = int(step_str)
        except Exception:
            # 理论上不会发生, 正则已经保证都是数字
            logger.warning("文件名中 step 部分无法转为 int, 跳过: %s", f.name)
            continue

        if step in step_to_file:
            raise RuntimeError(
                f"发现重复的 step={step} 对应多个文件: "
                f"{step_to_file[step]} 与 {f}"
            )
        step_to_file[step] = f

    if not step_to_file:
        raise RuntimeError(f"目录下未找到任何 eval_step_XXX.jsonl 文件: {exp_dir}")

    logger.info(
        "在目录 %s 中发现 %d 个 eval_step_XXX.jsonl 文件 (steps=%s)",
        exp_dir,
        len(step_to_file),
        sorted(step_to_file.keys()),
    )

    return step_to_file


# -----------------------------
# 工具函数: 调用 cal_score.py
# -----------------------------
def run_cal_score(
    python_exe: str,
    cal_score_script: Path,
    in_jsonl: Path,
    out_jsonl: Path,
) -> Path:
    """调用现有的 cal_score.py，对单个 JSONL 文件进行打分

    调用方式:
        python cal_score_script -i in_jsonl -o out_jsonl

    cal_score.py 默认会在 out_jsonl 后缀基础上加 ".metrics.json" 作为总体指标文件,
    即: metrics_path = out_jsonl + ".metrics.json"

    返回:
        metrics_path: Path
    """
    cmd = [
        python_exe,
        str(cal_score_script),
        "-i",
        str(in_jsonl),
        "-o",
        str(out_jsonl),
    ]

    logger.info("开始评估: in=%s  out=%s", in_jsonl, out_jsonl)
    try:
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        raise RuntimeError(
            f"调用 cal_score.py 失败: cmd={cmd}, returncode={e.returncode}"
        ) from e
    except Exception as e:
        raise RuntimeError(
            f"调用 cal_score.py 过程中出现异常: cmd={cmd}, error={e}"
        ) from e

    metrics_path = Path(str(out_jsonl) + ".metrics.json")
    if not metrics_path.is_file():
        raise RuntimeError(
            f"cal_score.py 执行完成但未找到 metrics 文件: {metrics_path}"
        )

    logger.info("评估完成, metrics: %s", metrics_path)
    return metrics_path


# -----------------------------
# 工具函数: 解析单个 metrics.json
# -----------------------------
def parse_metrics(metrics_path: Path) -> Dict[str, float]:
    """从 metrics JSON 中提取需要的 5 个 strict_accuracy 指标

    指标包括:
        - overall_strict_accuracy
        - judge_strict_accuracy
        - compare_strict_accuracy
        - diagnosis_strict_accuracy
        - choice_strict_accuracy
    若某个题型不存在, 则对应值为 0.0。
    """
    try:
        data = json.loads(metrics_path.read_text(encoding="utf-8"))
    except Exception as e:
        raise RuntimeError(
            f"读取或解析 metrics JSON 失败: path={metrics_path}, error={e}"
        ) from e

    overall = float(data.get("strict_accuracy", 0.0))
    per_type = data.get("per_type", {}) or {}

    def _get_type_acc(qtype: str) -> float:
        info = per_type.get(qtype, {}) or {}
        return float(info.get("strict_accuracy", 0.0))

    metrics = {
        "overall_strict_accuracy": overall,
        "judge_strict_accuracy": _get_type_acc("judge"),
        "compare_strict_accuracy": _get_type_acc("compare"),
        "diagnosis_strict_accuracy": _get_type_acc("diagnosis"),
        "choice_strict_accuracy": _get_type_acc("choice"),
    }

    return metrics

# -----------------------------
# 工具函数: 绘制折线图
# -----------------------------
def plot_metrics(
    step_metrics: Dict[int, Dict[str, float]],
    out_dir: Path,
) -> None:
    """根据 step->metrics 绘制折线图并保存

    调整后的策略:
        - 每个 strict_accuracy 指标单独画一张图, 避免所有曲线挤在一张图里难以阅读。
        - 图像统一保存在 out_dir 目录中。
        - 样式要求:
            * X 轴为 step (升序)
            * 灰色背景, 有网格
            * 一条折线, 使用方形点 (marker='s') 和稍微深一点的颜色
    """
    if not step_metrics:
        raise RuntimeError("plot_metrics 收到空的 step_metrics, 无法绘图")

    out_dir.mkdir(parents=True, exist_ok=True)

    steps_sorted = sorted(step_metrics.keys())

    # 准备每个指标的序列
    series_dict = {
        "overall_strict_accuracy": [
            step_metrics[s]["overall_strict_accuracy"] for s in steps_sorted
        ],
        "judge_strict_accuracy": [
            step_metrics[s]["judge_strict_accuracy"] for s in steps_sorted
        ],
        "compare_strict_accuracy": [
            step_metrics[s]["compare_strict_accuracy"] for s in steps_sorted
        ],
        "diagnosis_strict_accuracy": [
            step_metrics[s]["diagnosis_strict_accuracy"] for s in steps_sorted
        ],
        "choice_strict_accuracy": [
            step_metrics[s]["choice_strict_accuracy"] for s in steps_sorted
        ],
    }

    # 为了区分, 每个指标用一个不同但都比较深的颜色
    color_map = {
        "overall_strict_accuracy": "#1f77b4",   # 深蓝
        "judge_strict_accuracy": "#d62728",     # 深红
        "compare_strict_accuracy": "#2ca02c",   # 深绿
        "diagnosis_strict_accuracy": "#9467bd", # 深紫
        "choice_strict_accuracy": "#ff7f0e",    # 橙色
    }

    for metric_name, values in series_dict.items():
        plt.figure(figsize=(10, 6))
        ax = plt.gca()

        # 灰色背景 + 网格
        ax.set_facecolor("#f0f0f0")
        ax.grid(True, linestyle="--", alpha=0.5)

        color = color_map.get(metric_name, "#1f77b4")

        ax.plot(
            steps_sorted,
            values,
            marker="s",
            linestyle="-",
            linewidth=2,
            markersize=6,
            color=color,
        )

        ax.set_xlabel("step")
        ax.set_ylabel(metric_name)
        ax.set_title(f"{metric_name} vs Step")

        # 每个点标一个小方框就够了, legend 反而没什么必要, 因为一张图只有一条线
        plt.tight_layout()

        out_png = out_dir / f"{metric_name}.png"
        try:
            plt.savefig(out_png, dpi=150)
            logger.info("评估曲线已保存: %s", out_png)
        except Exception as e:
            raise RuntimeError(
                f"保存评估曲线失败: path={out_png}, error={e}"
            ) from e
        finally:
            plt.close()


# -----------------------------
# 主入口逻辑
# -----------------------------
def evaluate_steps(
    exp_dir: Path,
    python_exe: str,
    cal_score_script: Path,
    out_dir: Path,
) -> None:
    """对某个实验输出目录下的所有 eval_step_XXX.jsonl 进行批量评估

    参数:
        exp_dir:           包含 eval_step_XXX.jsonl 的目录
        python_exe:        Python 解释器路径 (一般用 sys.executable)
        cal_score_script:  单文件评估脚本 cal_score.py 的路径
        out_dir:           “目标路径”：单个 step 的打分结果(*.scored.jsonl 和对应 .metrics.json)输出目录
                           本函数会在 out_dir 下再创建一个 scored 子目录, 用于保存
                           总体 summary JSON 和所有折线图, 使目录结构更干净。
    """
    # step 级别的输出目录（每个 eval_step_* 的 scored.jsonl / .metrics.json 放这里）
    out_dir.mkdir(parents=True, exist_ok=True)

    # 汇总指标与图像的全局目录: out_dir/scored
    global_dir: Path = out_dir / "scored"
    global_dir.mkdir(parents=True, exist_ok=True)

    # 发现所有 step->文件
    step_to_file = discover_step_files(exp_dir)

    # 逐个 step 调用 cal_score.py, 顺序执行
    step_metrics: Dict[int, Dict[str, float]] = {}

    for step in sorted(step_to_file.keys()):
        in_jsonl = step_to_file[step]
        # 每个 step 的打分结果直接写到 out_dir 下
        out_jsonl = out_dir / f"eval_step_{step}.scored.jsonl"

        metrics_path = run_cal_score(
            python_exe=python_exe,
            cal_score_script=cal_score_script,
            in_jsonl=in_jsonl,
            out_jsonl=out_jsonl,
        )
        metrics = parse_metrics(metrics_path)
        step_metrics[step] = metrics

    # 写汇总 JSON (放在 global_dir)
    summary = {
        str(step): metrics
        for step, metrics in sorted(step_metrics.items(), key=lambda x: x[0])
    }
    summary_path = global_dir / "step_metrics_summary.json"
    try:
        summary_path.write_text(
            json.dumps(summary, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        logger.info("step 评估汇总已写入: %s", summary_path)
    except Exception as e:
        raise RuntimeError(
            f"写入 step_metrics_summary.json 失败: path={summary_path}, error={e}"
        ) from e

    # 画图 (各 metric 各一张图, 全部放在 global_dir 下)
    plot_metrics(step_metrics, out_dir=global_dir)


def parse_args() -> argparse.Namespace:
    """命令行参数解析"""
    parser = argparse.ArgumentParser(
        description="批量评估 eval_step_XXX.jsonl 并汇总 strict_accuracy 曲线",
    )
    parser.add_argument(
        "--exp-dir",
        "-id",
        type=str,
        required=True,
        help="评估输出目录, 内含 eval_step_XXX.jsonl (例如 experiments/eval_output/sub_exp_1_p20)",
    )
    parser.add_argument(
        "--cal-score-script",
        type=str,
        default="fundus_reasoner/eval/cal_score.py",
        help="单文件评估脚本 cal_score.py 的路径 (默认: fundus_reasoner/eval/cal_score.py)",
    )
    parser.add_argument(
        "--out-dir",
        "-o",
        type=str,
        default=None,
        help="批量评估结果输出目录 (默认: <exp-dir>/scored)",
    )
    parser.add_argument(
        "--python-exe",
        type=str,
        default=sys.executable,
        help="调用 cal_score.py 使用的 Python 解释器路径 (默认: 当前解释器)",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    exp_dir = Path(args.exp_dir)
    cal_score_script = Path(args.cal_score_script)

    if args.out_dir is not None:
        out_dir = Path(args.out_dir)
    else:
        out_dir = exp_dir / "scored"

    evaluate_steps(
        exp_dir=exp_dir,
        python_exe=args.python_exe,
        cal_score_script=cal_score_script,
        out_dir=out_dir,
    )


if __name__ == "__main__":
    main()
